/*
 * Copyright 2023-2025, Haiku, Inc. All rights reserved.
 * Distributed under the terms of the MIT License.
 *
 * 龙芯LoongArch64架构TLB管理实现
 * 本文件实现LoongArch64架构的TLB管理功能
 */

#include <KernelExport.h>

#include <arch/cpu.h>
#include <arch/loongarch64/arch_mmu.h>
#include <boot/kernel_args.h>
#include <kernel.h>
#include <vm/vm.h>
#include <vm/VMAddressSpace.h>

#include "LoongArchVMTranslationMap.h"


//#define TRACE_LOONGARCH_TLB
#ifdef TRACE_LOONGARCH_TLB
#   define TRACE(x...) dprintf("loongarch_tlb: " x)
#else
#   define TRACE(x...) ;
#endif


// 定义LoongArch64 TLB相关CSR寄存器
#define CSR_TLBIDX    0x10    // TLB索引
#define CSR_TLBEHI    0x11    // TLB入口高位
#define CSR_TLBELO0   0x12    // TLB入口低位0
#define CSR_TLBELO1   0x13    // TLB入口低位1
#define CSR_ASID      0x18    // 地址空间标识符
#define CSR_PGDL      0x19    // 页表基址低位
#define CSR_PGDH      0x1a    // 页表基址高位
#define CSR_PGD       0x1b    // 页表基址

// 定义TLB操作指令
#define TLBSRCH     0x0       // 搜索TLB
#define TLBRD       0x1       // 读取TLB
#define TLBWR       0x2       // 写入TLB
#define TLBFILL     0x3       // 填充TLB
#define TLBCLR      0x4       // 清除TLB
#define TLBFLUSH    0x5       // 刷新TLB
#define INVTLB      0x6       // 无效化TLB

// 定义TLBIDX寄存器位
#define TLBIDX_PS_SHIFT  24    // 页大小偏移
#define TLBIDX_PS_MASK   0x3f  // 页大小掩码
#define TLBIDX_INDEX_SHIFT 0   // 索引偏移
#define TLBIDX_INDEX_MASK 0xfff // 索引掩码

// 定义INVTLB操作码
#define INVTLB_ALL          0    // 无效化所有TLB
#define INVTLB_ASID         1    // 无效化指定ASID的TLB
#define INVTLB_VA           2    // 无效化指定VA的TLB
#define INVTLB_ASID_VA      3    // 无效化指定ASID和VA的TLB
#define INVTLB_ASID_VA_PAGE 4    // 无效化指定ASID和VA页的TLB

// 定义页大小
#define PS_4K   12   // 4KB页
#define PS_8K   13   // 8KB页
#define PS_16K  14   // 16KB页
#define PS_32K  15   // 32KB页
#define PS_64K  16   // 64KB页
#define PS_128K 17   // 128KB页
#define PS_256K 18   // 256KB页
#define PS_512K 19   // 512KB页
#define PS_1M   20   // 1MB页
#define PS_2M   21   // 2MB页
#define PS_4M   22   // 4MB页
#define PS_8M   23   // 8MB页
#define PS_16M  24   // 16MB页
#define PS_32M  25   // 32MB页
#define PS_64M  26   // 64MB页
#define PS_128M 27   // 128MB页
#define PS_256M 28   // 256MB页
#define PS_512M 29   // 512MB页
#define PS_1G   30   // 1GB页
#define PS_2G   31   // 2GB页
#define PS_4G   32   // 4GB页

// 全局变量
static uint32 sTlbEntries = 0;  // TLB条目数量
static uint32 sTlbPageSizes = 0; // 支持的页大小


// 读取CSR寄存器
static inline uint64
read_csr(uint32 reg)
{
    uint64 val;
    asm volatile("csrrd %0, %1" : "=r"(val) : "i"(reg));
    return val;
}

// 写入CSR寄存器
static inline void
write_csr(uint32 reg, uint64 val)
{
    asm volatile("csrwr %0, %1" : : "r"(val), "i"(reg));
}

// 执行TLB操作
static inline void
exec_tlb_op(uint32 op)
{
    asm volatile("tlb %0" : : "i"(op) : "memory");
}

// 执行INVTLB操作
static inline void
exec_invtlb(uint32 op, uint64 asid, uint64 va)
{
    asm volatile("invtlb %0, %1, %2" : : "i"(op), "r"(asid), "r"(va) : "memory");
}


// 初始化TLB
status_t
loongarch_tlb_init(kernel_args *args)
{
    TRACE("Initializing TLB\n");

    // 检测TLB条目数量
    // 通过读取CPUCFG寄存器获取TLB信息
    uint32 cpucfg2 = read_cpucfg(2);
    sTlbEntries = ((cpucfg2 >> 16) & 0x3ff) + 1;
    
    // 检测支持的页大小
    uint32 cpucfg3 = read_cpucfg(3);
    sTlbPageSizes = (cpucfg3 >> 8) & 0xff;

    TRACE("TLB entries: %u, supported page sizes: 0x%x\n", sTlbEntries, sTlbPageSizes);

    return B_OK;
}

// 初始化每个CPU的TLB
status_t
loongarch_tlb_init_percpu(kernel_args *args, int cpu)
{
    TRACE("Initializing TLB for CPU %d\n", cpu);

    // 设置页表基址
    uint64 pgd = args->arch_args.phys_pgd;
    write_csr(CSR_PGD, pgd);

    // 清除TLB
    exec_invtlb(INVTLB_ALL, 0, 0);

    return B_OK;
}

// 无效化TLB范围
void
loongarch_invalidate_tlb_range(addr_t start, addr_t end)
{
    // 确保地址对齐到页边界
    start = ROUNDDOWN(start, B_PAGE_SIZE);
    end = ROUNDUP(end, B_PAGE_SIZE);

    // 遍历范围内的每个页面，无效化TLB
    for (addr_t addr = start; addr < end; addr += B_PAGE_SIZE) {
        exec_invtlb(INVTLB_VA, 0, addr);
    }
}

// 无效化TLB列表
void
loongarch_invalidate_tlb_list(addr_t pages[], int num_pages)
{
    for (int i = 0; i < num_pages; i++) {
        exec_invtlb(INVTLB_VA, 0, pages[i]);
    }
}

// 全局无效化TLB
void
loongarch_global_tlb_invalidate(void)
{
    exec_invtlb(INVTLB_ALL, 0, 0);
}

// 用户空间无效化TLB
void
loongarch_user_tlb_invalidate(void)
{
    // 获取当前ASID
    uint32 asid = read_csr(CSR_ASID);
    
    // 无效化指定ASID的所有TLB条目
    exec_invtlb(INVTLB_ASID, asid, 0);
}

// 手动填充TLB条目
status_t
loongarch_fill_tlb(addr_t va, phys_addr_t pa, uint32 attributes, uint32 memoryType)
{
    // 构建TLB条目
    uint64 tlbelo0 = (pa & kPtePhysAddrMask);
    
    // 设置属性标志
    if ((attributes & B_READ_AREA) != 0)
        tlbelo0 |= kPteRead;
    if ((attributes & B_WRITE_AREA) != 0)
        tlbelo0 |= kPteWrite;
    if ((attributes & B_EXECUTE_AREA) != 0)
        tlbelo0 |= kPteExec;
    
    // 添加全局页标志和有效位
    tlbelo0 |= kPteGlobal | kPteValid;
    
    // 设置TLB寄存器
    write_csr(CSR_TLBEHI, va);
    write_csr(CSR_TLBELO0, tlbelo0);
    write_csr(CSR_TLBELO1, 0); // 不使用第二个页面
    
    // 设置页大小为4KB
    uint64 tlbidx = read_csr(CSR_TLBIDX);
    tlbidx &= ~(TLBIDX_PS_MASK << TLBIDX_PS_SHIFT);
    tlbidx |= (PS_4K << TLBIDX_PS_SHIFT);
    write_csr(CSR_TLBIDX, tlbidx);
    
    // 执行TLB填充操作
    exec_tlb_op(TLBFILL);
    
    return B_OK;
}